import json

from stateSecret.sm4_1 import CryptSM4, SM4_ENCRYPT, SM4_DECRYPT
import base64


class SM4Utils:

    def __init__(self, ):
        self.crypt_sm4 = CryptSM4()

    # 加密方法
    def encryptData_ECB(self, origin_text: str, secret_key: str) -> str:
        """
        sm4加密
        :param origin_text: 原始数据
        :param secret_key: 秘钥
        :return: 加密后的信息
        """
        origin_text = bytes(str(origin_text), encoding='utf8')
        secret_key = bytes(str(secret_key), encoding='utf8')
        self.crypt_sm4.set_key(secret_key, SM4_ENCRYPT)
        encrypt_value = self.crypt_sm4.crypt_ecb(origin_text)
        return base64.b64encode(encrypt_value).decode('utf-8', 'ignore')

    def decryptData_ECB(self, secret_text: str, secret_key: str) -> str:
        """
        sm4解密
        :param secret_text: 加密后的信息
        :param secret_key: 秘钥
        :return: 解密后的信息
        """
        secret_key = bytes(str(secret_key), encoding='utf8')
        self.crypt_sm4.set_key(secret_key, SM4_DECRYPT)
        byt_cipher_text = base64.b64decode(secret_text)
        decrypt_value = self.crypt_sm4.crypt_ecb(byt_cipher_text)
        return decrypt_value.decode('utf-8', 'ignore')


SM4_Utils = SM4Utils()
if __name__ == '__main__':
    sec = "JeF38U9wT9wlMfs2"
    plain_text = {"t": 1625451188275, "userToken": "c49bdd6a558948ec86c957c8bd07c1c8"}
    sec_text = SM4_Utils.encryptData_ECB(str(plain_text), sec)
    print(sec_text)
    print(SM4_Utils.decryptData_ECB(sec_text, sec))
